#include "yolo_v6.h"
#include "../txr_algo_dlm_rec_v1/common/common.hpp"

#include "yaml-cpp/yaml.h"

#include "boost/process.hpp"
#include "boost/process/windows.hpp"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/path.hpp"
#include "boost/filesystem/operations.hpp"
#include "boost/filesystem/fstream.hpp"

#include "dll.h"

using namespace CryptoPP;

std::vector<char> decryptFile(const std::vector<char> & v_encrypt, const std::string& key, const std::string& iv)
{
	std::vector<char> v_decrypt;
	try
	{
		CBC_Mode<AES>::Decryption decryption((const byte*)key.data(), key.size(), (const byte*)iv.data());

		std::vector<char> _decrypt = v_encrypt;
		decryption.ProcessData(
			reinterpret_cast<CryptoPP::byte*>(_decrypt.data()),
			reinterpret_cast<const CryptoPP::byte*>(v_encrypt.data()),
			v_encrypt.size());
		
		byte paddingByte = _decrypt.back();
		size_t paddingSize = static_cast<size_t>(paddingByte);
		_decrypt.resize(v_encrypt.size() - paddingSize);

		v_decrypt = _decrypt;
	}
	catch (const CryptoPP::Exception& e)
	{
		std::cerr << e.what() << std::endl;
	}
	return v_decrypt;
}


yolo_v6::yolo_v6()
{

}
yolo_v6::~yolo_v6()
{

}

std::vector<char> get_file_data(std::string file_path)
{
	std::ifstream file(file_path, std::ios_base::binary);
	assert(file.good());
	file.seekg(0, std::ios_base::end);
	int file_len = file.tellg();
	file.seekg(0, std::ios_base::beg);

	std::vector<char> data(file_len, 0);
	file.read(data.data(), file_len);

	file.close();
	return data;
}

void write_file_data(std::vector<char> data, std::string file_path)
{
	std::ofstream file(file_path, std::ios_base::binary);
	file.write(data.data(), data.size());
	file.close();
}

void encrypt_data(std::vector<char> data)
{
	char trans = 55;
	for (int i = 0; i < data.size(); ++i)
	{
		data[i] += trans;
		trans += 11;
	}
}

void encryptFile(st_encrypt_info info, std::vector<char> engine_data, std::string target_path)
{
	info.engine_size = engine_data.size();
	std::vector<char> target_data(sizeof(info), 0);
	*((st_encrypt_info*)(target_data.data())) = info;

	target_data.insert(target_data.end(), engine_data.begin(), engine_data.end());

	encrypt_data(target_data);
	write_file_data(target_data, target_path);
	std::cout << "writing encrypt file..." << target_path <<" done!" << std::endl;

}


bool yolo_v6::TransModel(st_trans_model_input input)
{
	namespace fs = boost::filesystem;
	YAML::Node root = YAML::LoadFile(input.cfg_path);
	YAML::Node config = root["REC"];
	
	//std::string onnx_file = config["onnx_file"].as<std::string>();
	//std::string engine_file = config["engine_file"].as<std::string>();
	int batch_size = config["BATCH_SIZE"].as<int>();
	int input_channel = config["INPUT_CHANNEL"].as<int>();
	int image_width = config["IMAGE_WIDTH"].as<int>();
	int image_height = config["IMAGE_HEIGHT"].as<int>();
	bool dynamic = config["Dynamic"].as<bool>();

	std::vector<char> v_model, v_engine;

	fs::path p(input.model_path.c_str());
	if (p.extension().string() == ".onnx")
	{
		v_model = get_file_data(p.string());
	}
	else if (p.extension().string() == ".ennx")
	{
		std::vector<char> v_e_model;
		v_e_model = get_file_data(p.string());
		std::string key = "irkdh_573%3?iq5h";
		std::string iv = "1234567890123456";
		v_model = decryptFile(v_e_model, key, iv);
	}
	else
	{
		return false;
	}
	{
		nvinfer1::ICudaEngine* engine = nullptr;
		onnxToTRTModel(v_model, v_engine, engine, batch_size, image_width, image_height, dynamic);
		assert(engine != nullptr);
		if (engine == nullptr)
		{
			return false;
		}
		engine->destroy();
	}

	st_encrypt_info info;
	sprintf_s(info.gpu_model, 100, "%s", input.gpu_model.c_str());
	info.batch_size = batch_size;
	info.input_channel = input_channel;
	info.image_width = image_width;
	info.image_height = image_height;
	

	std::string encrypt_file = input.trans_path;
	encryptFile(info, v_engine, encrypt_file);

	return true;
}